package com.youlu.campus.starter.trace.thread;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 支持日志追踪 的 线程池
 *
 * @Description: TraceLogThreadPoolExecutor
 * @Date 2025/8/29
 * @module TraceLogThreadPoolExecutor
 * @author: su
 */
public class LogTraceThreadPoolExecutor extends ThreadPoolExecutor {


    public LogTraceThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public LogTraceThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public LogTraceThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public LogTraceThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public Future<?> submit(Runnable task) {
        // 1. 在提交任务时，获取当前线程（调用者线程）的MDC上下文副本
        Map<String, String> context = MDC.getCopyOfContextMap();
        // 2. 包装原始任务，使其在执行前设置MDC上下文，执行后清理
        Runnable wrappedTask = wrap(task, context);
        // 3. 将包装后的任务提交给父类线程池执行
        return super.submit(wrappedTask);
    }

    /**
     * 包装Runnable任务，用于传递MDC上下文
     */
    private Runnable wrap(Runnable task, Map<String, String> context) {
        return () -> {
            // 任务执行前，将父线程的MDC上下文设置到当前子线程中
            if (context != null) {
                MDC.setContextMap(context);
            }
            try {
                // 执行原始任务的run方法
                task.run();
            } finally {
                // 确保执行后清理子线程的MDC，避免内存泄漏和上下文污染
                MDC.clear();
            }
        };
    }
}
